from typing import Any
from transformers import BertTokenizer, BertModel
from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct

from file import read_json_files_to_array

tokenizer = BertTokenizer.from_pretrained("D:/bert-base-chinese")
model = BertModel.from_pretrained("D:/bert-base-chinese")


client = QdrantClient(host="10.254.32.20", port=6333)

# qdrantKey = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.Yv7GXXOZw37Ci0v1mtnxoaEar9DFSXcA4VwprhTLPBk"
# client = QdrantClient(
#     host="https://1298c697-e63d-462a-bc7b-faf3c988bca8.us-west-1-0.aws.cloud.qdrant.io:6333",
#     api_key=qdrantKey,
# )

collection_name = "webui"

# 创建
# client.create_collection(
#     collection_name=collection_name,
#     # vectors_config=VectorParams(size=4, distance=Distance.DOT),
#     vectors_config=VectorParams(size=768, distance=Distance.DOT),
# )


def EmbeddingText(text: str) -> Any:
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    return outputs.last_hidden_state[:, 0, :].tolist()[0]


# ret_query = client.query_points(
#     collection_name="restfullapi", query=vector, score_threshold=0.6, limit=30
# )


# operation_info = client.upsert(
#     collection_name=collection_name,
#     wait=True,
#     points=[
#         PointStruct(id=1, vector=[0.05, 0.61, 0.76, 0.74], payload={"city": "Berlin"}),
#         PointStruct(id=2, vector=[0.19, 0.81, 0.75, 0.11], payload={"city": "London"}),
#         PointStruct(id=3, vector=[0.36, 0.55, 0.47, 0.94], payload={"city": "Moscow"}),
#         PointStruct(
#             id=4, vector=[0.18, 0.01, 0.85, 0.80], payload={"city": "New York"}
#         ),
#         PointStruct(id=5, vector=[0.24, 0.18, 0.22, 0.44], payload={"city": "Beijing"}),
#         PointStruct(id=6, vector=[0.35, 0.08, 0.11, 0.44], payload={"city": "Mumbai"}),
#     ],
# )

if __name__ == "__main__":
    files = read_json_files_to_array("data")
    for item in files:
        print("item", item)
        text = "\n".join(item["content"])
        vector = EmbeddingText(text)
        client.upsert(
            collection_name=collection_name,
            wait=True,
            points=[PointStruct(id=item["id"], vector=vector, payload=item["data"])],
        )
